"""
数据库操作工具类
提供知识库向量检索、CRUD操作等功能
"""

from typing import List, Dict, Any, Optional, Tuple
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, desc, func, update, delete
from sqlalchemy.orm import selectinload
import numpy as np
import uuid
import json
from datetime import datetime, timedelta, timezone
import logging
from utils.logger import logger

from app.core.database_core import (
    KnowledgeBase,
    AgentTask,
    UserInteraction
)

# 对于缺失的模型，创建占位符
try:
    # 尝试从database_core导入其余模型
    from app.core.database_core import VECTOR_TYPE_AVAILABLE, Vector, WorkflowState, TaskState
except ImportError:
    logger.warning("部分数据库模型未找到，将使用默认值")
    VECTOR_TYPE_AVAILABLE = False
    # 创建简单的占位符类
    class Vector:
        pass
    class WorkflowState:
        pass
    class TaskState:
        pass


class DatabaseUtils:
    """数据库操作工具类"""
    
    @staticmethod
    async def add_knowledge(session: AsyncSession, content: str, embedding: List[float], metadata_: Optional[Dict] = None) -> KnowledgeBase:
        """添加知识到知识库"""
        try:
            # 验证embedding维度
            if len(embedding) != 1536:
                logger.warning(f"嵌入向量维度不匹配，预期1536，实际{len(embedding)}")
            
            knowledge = KnowledgeBase(
                content=content,
                embedding=embedding,
                metadata_=metadata_ or {}
            )
            session.add(knowledge)
            await session.commit()
            await session.refresh(knowledge)
            logger.info(f"成功添加知识条目，ID: {knowledge.id}")
            return knowledge
        except Exception as e:
            await session.rollback()
            logger.error(f"添加知识失败: {str(e)}")
            raise
    
    @staticmethod
    async def search_knowledge_by_vector(
        session: AsyncSession, 
        query_embedding: List[float], 
        top_k: int = 5,
        filter_metadata: Optional[Dict] = None
    ) -> List[Dict[str, Any]]:
        """通过向量相似度搜索知识库
        
        Args:
            session: 数据库会话
            query_embedding: 查询向量
            top_k: 返回前k个结果
            filter_metadata: 元数据过滤条件
            
        Returns:
            包含知识条目和相似度的列表
        """
        if not VECTOR_TYPE_AVAILABLE:
            logger.warning("向量检索功能不可用，返回空结果")
            return []
        
        try:
            # 构建基础查询
            stmt = select(
                KnowledgeBase,
                (KnowledgeBase.embedding.cosine_distance(query_embedding)).label('distance')
            )
            
            # 添加元数据过滤
            if filter_metadata:
                for key, value in filter_metadata.items():
                    stmt = stmt.where(KnowledgeBase.metadata_[key] == value)
            
            # 按距离排序并限制结果数
            stmt = stmt.order_by('distance').limit(top_k)
            
            result = await session.execute(stmt)
            items = result.all()
            
            # 格式化结果
            search_results = []
            for knowledge, distance in items:
                # cosine_distance越小越相似，转换为相似度分数
                similarity = 1 - float(distance) if distance is not None else 0.0
                search_results.append({
                    'id': knowledge.id,
                    'content': knowledge.content,
                    'embedding': knowledge.embedding,
                    'metadata': knowledge.metadata_,
                    'similarity': similarity,
                    'created_at': knowledge.created_at,
                    'updated_at': knowledge.updated_at
                })
            
            logger.info(f"向量检索完成，找到{len(search_results)}条匹配结果")
            return search_results
            
        except Exception as e:
            logger.error(f"向量检索失败: {str(e)}")
            return []
    
    @staticmethod
    async def batch_add_knowledge(
        session: AsyncSession,
        knowledge_items: List[Dict[str, Any]]
    ) -> List[int]:
        """批量添加知识库条目
        
        Args:
            session: 数据库会话
            knowledge_items: 知识条目列表，每个条目包含content和embedding
            
        Returns:
            成功添加的知识ID列表
        """
        try:
            added_ids = []
            
            # 每批处理100条，避免事务过大
            batch_size = 100
            for i in range(0, len(knowledge_items), batch_size):
                batch = knowledge_items[i:i + batch_size]
                knowledge_objects = []
                
                for item in batch:
                    knowledge = KnowledgeBase(
                        content=item['content'],
                        embedding=item['embedding'],
                        metadata_=item.get('metadata', {})
                    )
                    knowledge_objects.append(knowledge)
                    
                session.add_all(knowledge_objects)
                await session.flush()  # 获取ID但不提交事务
                
                # 收集ID
                added_ids.extend([k.id for k in knowledge_objects])
                logger.info(f"批量添加进度: {min(i + batch_size, len(knowledge_items))}/{len(knowledge_items)}")
            
            await session.commit()
            logger.info(f"批量添加完成，共添加{len(added_ids)}条知识")
            return added_ids
            
        except Exception as e:
            await session.rollback()
            logger.error(f"批量添加失败: {str(e)}")
            raise
    
    @staticmethod
    async def create_agent_task(
        session: AsyncSession,
        task_type: str,
        input_data: Dict[str, Any],
        agent_id: str,
        priority: int = 0
    ) -> AgentTask:
        """创建智能体任务
        
        Args:
            session: 数据库会话
            task_type: 任务类型
            input_data: 输入数据
            agent_id: 智能体ID
            priority: 优先级，数字越大优先级越高
            
        Returns:
            创建的任务对象
        """
        try:
            task = AgentTask(
                id=str(uuid.uuid4()),
                task_type=task_type,
                status='pending',
                input_data=input_data,
                agent_id=agent_id,
                priority=priority
            )
            session.add(task)
            await session.commit()
            await session.refresh(task)
            logger.info(f"创建任务成功: {task.id}, 类型: {task_type}")
            return task
        except Exception as e:
            await session.rollback()
            logger.error(f"创建任务失败: {str(e)}")
            raise
    
    @staticmethod
    async def update_task_status(
        session: AsyncSession,
        task_id: str,
        status: str,
        output_data: Optional[Dict[str, Any]] = None
    ) -> bool:
        """更新任务状态
        
        Args:
            session: 数据库会话
            task_id: 任务ID
            status: 新状态
            output_data: 输出数据（如果任务完成）
            
        Returns:
            是否更新成功
        """
        try:
            result = await session.execute(
                select(AgentTask).where(AgentTask.id == task_id)
            )
            task = result.scalar_one_or_none()
            
            if not task:
                logger.warning(f"任务不存在: {task_id}")
                return False
            
            task.status = status
            if output_data is not None:
                task.output_data = output_data
            
            # 如果任务完成，记录完成时间
            if status in ['completed', 'failed']:
                from datetime import datetime, timezone
                task.completed_at = datetime.now(timezone.utc)
            
            await session.commit()
            logger.info(f"任务状态更新成功: {task_id} -> {status}")
            return True
            
        except Exception as e:
            await session.rollback()
            logger.error(f"更新任务状态失败: {str(e)}")
            return False
    
    @staticmethod
    async def save_user_interaction(
        session: AsyncSession,
        user_id: str,
        query: str,
        response: str,
        session_id: str,
        query_embedding: Optional[List[float]] = None
    ) -> UserInteraction:
        """保存用户交互记录
        
        Args:
            session: 数据库会话
            user_id: 用户ID
            query: 用户查询
            response: 系统响应
            session_id: 会话ID
            query_embedding: 查询的向量表示
            
        Returns:
            创建的交互记录
        """
        try:
            interaction = UserInteraction(
                user_id=user_id,
                query=query,
                response=response,
                session_id=session_id,
                query_embedding=query_embedding
            )
            session.add(interaction)
            await session.commit()
            await session.refresh(interaction)
            logger.info(f"保存用户交互记录成功: {interaction.id}")
            return interaction
        except Exception as e:
            await session.rollback()
            logger.error(f"保存用户交互记录失败: {str(e)}")
            raise
    
    @staticmethod
    async def search_similar_queries(
        session: AsyncSession,
        query_embedding: List[float],
        top_k: int = 3,
        min_similarity: float = 0.7
    ) -> List[Dict[str, Any]]:
        """搜索相似的历史查询
        
        Args:
            session: 数据库会话
            query_embedding: 查询向量
            top_k: 返回前k个结果
            min_similarity: 最小相似度阈值
            
        Returns:
            相似查询列表
        """
        if not VECTOR_TYPE_AVAILABLE:
            return []
        
        try:
            stmt = select(
                UserInteraction,
                (UserInteraction.query_embedding.cosine_distance(query_embedding)).label('distance')
            ).where(
                UserInteraction.query_embedding.isnot(None)
            ).order_by('distance').limit(top_k)
            
            result = await session.execute(stmt)
            items = result.all()
            
            similar_queries = []
            for interaction, distance in items:
                similarity = 1 - float(distance) if distance is not None else 0.0
                if similarity >= min_similarity:
                    similar_queries.append({
                        'id': interaction.id,
                        'user_id': interaction.user_id,
                        'query': interaction.query,
                        'response': interaction.response,
                        'similarity': similarity,
                        'created_at': interaction.created_at
                    })
            
            return similar_queries
            
        except Exception as e:
            logger.error(f"搜索相似查询失败: {str(e)}")
            return []
    
    @staticmethod
    async def delete_old_knowledge(
        session: AsyncSession,
        days: int = 90
    ) -> int:
        """删除旧的知识库条目
        
        Args:
            session: 数据库会话
            days: 删除多少天前的条目
            
        Returns:
            删除的条目数
        """
        try:
            from datetime import datetime, timedelta, timezone
            cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
            
            # 先查询要删除的条目数量
            result = await session.execute(
                select(func.count(KnowledgeBase.id)).where(
                    KnowledgeBase.created_at < cutoff_date
                )
            )
            count = result.scalar()
            
            if count > 0:
                # 执行删除
                from sqlalchemy import delete
                stmt = delete(KnowledgeBase).where(
                    KnowledgeBase.created_at < cutoff_date
                )
                await session.execute(stmt)
                await session.commit()
                logger.info(f"删除了{count}条{days}天前的知识库条目")
            
            return count
            
        except Exception as e:
            await session.rollback()
            logger.error(f"删除旧知识库条目失败: {str(e)}")
            return 0


    @staticmethod
    async def save_workflow_state(session: AsyncSession, workflow_data: Dict[str, Any]) -> str:
        """保存工作流状态到数据库
        
        Args:
            session: 数据库会话
            workflow_data: 工作流数据
            
        Returns:
            工作流ID
        """
        try:
            # 转换复杂数据类型为JSON字符串
            workflow_state = WorkflowState(
                workflow_id=workflow_data["workflow_id"],
                task_type=workflow_data["task_type"],
                query=workflow_data["query"],
                results=json.dumps(workflow_data.get("results", {})),
                errors=json.dumps(workflow_data.get("errors", [])),
                agent_status=json.dumps(workflow_data.get("agent_status", {})),
                completed_tasks=json.dumps(workflow_data.get("completed_tasks", [])),
                metadata=json.dumps(workflow_data.get("metadata", {})),
                timestamp=datetime.fromisoformat(workflow_data.get("timestamp", datetime.now().isoformat()))
            )
            
            session.add(workflow_state)
            await session.commit()
            await session.refresh(workflow_state)
            
            # 保存任务状态
            if "completed_tasks" in workflow_data:
                for task_id in workflow_data["completed_tasks"]:
                    task_state = TaskState(
                        workflow_id=workflow_data["workflow_id"],
                        task_id=task_id,
                        status="completed",
                        timestamp=datetime.now(timezone.utc)
                    )
                    session.add(task_state)
            
            await session.commit()
            logger.info(f"工作流状态保存成功: {workflow_data['workflow_id']}")
            return workflow_data["workflow_id"]
            
        except Exception as e:
            await session.rollback()
            logger.error(f"保存工作流状态失败: {str(e)}")
            raise
    
    @staticmethod
    async def get_workflow_state(session: AsyncSession, workflow_id: str) -> Optional[Dict[str, Any]]:
        """从数据库获取工作流状态
        
        Args:
            session: 数据库会话
            workflow_id: 工作流ID
            
        Returns:
            工作流状态，如果不存在则返回None
        """
        try:
            stmt = select(WorkflowState).where(WorkflowState.workflow_id == workflow_id)
            result = await session.execute(stmt)
            workflow_state = result.scalars().first()
            
            if not workflow_state:
                return None
            
            # 转换JSON字符串为Python对象
            return {
                "id": workflow_state.id,
                "workflow_id": workflow_state.workflow_id,
                "task_type": workflow_state.task_type,
                "query": workflow_state.query,
                "results": json.loads(workflow_state.results),
                "errors": json.loads(workflow_state.errors),
                "agent_status": json.loads(workflow_state.agent_status),
                "completed_tasks": json.loads(workflow_state.completed_tasks),
                "metadata": json.loads(workflow_state.metadata),
                "timestamp": workflow_state.timestamp.isoformat(),
                "created_at": workflow_state.created_at.isoformat(),
                "updated_at": workflow_state.updated_at.isoformat()
            }
            
        except Exception as e:
            logger.error(f"获取工作流状态失败: {str(e)}")
            raise
    
    @staticmethod
    async def update_workflow_state(session: AsyncSession, workflow_id: str, updates: Dict[str, Any]) -> bool:
        """更新工作流状态
        
        Args:
            session: 数据库会话
            workflow_id: 工作流ID
            updates: 更新的字段
            
        Returns:
            是否更新成功
        """
        try:
            # 准备更新数据，转换复杂类型为JSON
            update_data = {}
            for key, value in updates.items():
                if isinstance(value, (dict, list)):
                    update_data[key] = json.dumps(value)
                else:
                    update_data[key] = value
            
            update_data["updated_at"] = datetime.now(timezone.utc)
            
            stmt = update(WorkflowState).where(
                WorkflowState.workflow_id == workflow_id
            ).values(**update_data)
            
            result = await session.execute(stmt)
            await session.commit()
            
            success = result.rowcount > 0
            logger.info(f"更新工作流状态 {'成功' if success else '失败'}: {workflow_id}")
            return success
            
        except Exception as e:
            await session.rollback()
            logger.error(f"更新工作流状态失败: {str(e)}")
            raise
    
    @staticmethod
    async def list_workflow_states(session: AsyncSession, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
        """列出工作流状态
        
        Args:
            session: 数据库会话
            limit: 返回数量限制
            offset: 偏移量
            
        Returns:
            工作流状态列表
        """
        try:
            stmt = select(WorkflowState).order_by(
                WorkflowState.created_at.desc()
            ).limit(limit).offset(offset)
            
            result = await session.execute(stmt)
            workflow_states = result.scalars().all()
            
            # 转换为字典列表
            return [
                {
                    "id": state.id,
                    "workflow_id": state.workflow_id,
                    "task_type": state.task_type,
                    "query": state.query,
                    "results": json.loads(state.results),
                    "errors": json.loads(state.errors),
                    "agent_status": json.loads(state.agent_status),
                    "completed_tasks": json.loads(state.completed_tasks),
                    "metadata": json.loads(state.metadata),
                    "timestamp": state.timestamp.isoformat(),
                    "created_at": state.created_at.isoformat(),
                    "updated_at": state.updated_at.isoformat()
                }
                for state in workflow_states
            ]
            
        except Exception as e:
            logger.error(f"列出工作流状态失败: {str(e)}")
            raise
    
    @staticmethod
    async def delete_workflow_state(session: AsyncSession, workflow_id: str) -> bool:
        """删除工作流状态
        
        Args:
            session: 数据库会话
            workflow_id: 工作流ID
            
        Returns:
            是否删除成功
        """
        try:
            # 先删除相关的任务状态
            task_stmt = delete(TaskState).where(TaskState.workflow_id == workflow_id)
            await session.execute(task_stmt)
            
            # 再删除工作流状态
            workflow_stmt = delete(WorkflowState).where(WorkflowState.workflow_id == workflow_id)
            result = await session.execute(workflow_stmt)
            
            await session.commit()
            
            success = result.rowcount > 0
            logger.info(f"删除工作流状态 {'成功' if success else '失败'}: {workflow_id}")
            return success
            
        except Exception as e:
            await session.rollback()
            logger.error(f"删除工作流状态失败: {str(e)}")
            raise
    
    @staticmethod
    async def save_task_state(session: AsyncSession, workflow_id: str, task_id: str, status: str, 
                            metadata: Optional[Dict[str, Any]] = None) -> str:
        """保存任务状态
        
        Args:
            session: 数据库会话
            workflow_id: 工作流ID
            task_id: 任务ID
            status: 任务状态
            metadata: 任务元数据
            
        Returns:
            任务状态ID
        """
        try:
            task_state = TaskState(
                workflow_id=workflow_id,
                task_id=task_id,
                status=status,
                metadata=json.dumps(metadata or {}),
                timestamp=datetime.now(timezone.utc)
            )
            
            session.add(task_state)
            await session.commit()
            await session.refresh(task_state)
            
            logger.info(f"任务状态保存成功: workflow={workflow_id}, task={task_id}, status={status}")
            return str(task_state.id)
            
        except Exception as e:
            await session.rollback()
            logger.error(f"保存任务状态失败: {str(e)}")
            raise
    
    @staticmethod
    async def get_workflow_tasks(session: AsyncSession, workflow_id: str) -> List[Dict[str, Any]]:
        """获取工作流的所有任务状态
        
        Args:
            session: 数据库会话
            workflow_id: 工作流ID
            
        Returns:
            任务状态列表
        """
        try:
            stmt = select(TaskState).where(
                TaskState.workflow_id == workflow_id
            ).order_by(TaskState.timestamp)
            
            result = await session.execute(stmt)
            task_states = result.scalars().all()
            
            return [
                {
                    "id": state.id,
                    "workflow_id": state.workflow_id,
                    "task_id": state.task_id,
                    "status": state.status,
                    "metadata": json.loads(state.metadata),
                    "timestamp": state.timestamp.isoformat(),
                    "created_at": state.created_at.isoformat()
                }
                for state in task_states
            ]
            
        except Exception as e:
            logger.error(f"获取工作流任务状态失败: {str(e)}")
            raise
    
    @staticmethod
    async def cleanup_old_workflows(session: AsyncSession, days: int = 7) -> int:
        """清理旧的工作流数据
        
        Args:
            session: 数据库会话
            days: 保留天数
            
        Returns:
            删除的工作流数量
        """
        try:
            cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
            
            # 先查询要删除的工作流ID
            stmt = select(WorkflowState.workflow_id).where(
                WorkflowState.created_at < cutoff_date
            )
            result = await session.execute(stmt)
            workflow_ids = result.scalars().all()
            
            if not workflow_ids:
                logger.info("没有需要清理的旧工作流")
                return 0
            
            # 删除相关的任务状态
            task_stmt = delete(TaskState).where(
                TaskState.workflow_id.in_(workflow_ids)
            )
            await session.execute(task_stmt)
            
            # 删除工作流状态
            workflow_stmt = delete(WorkflowState).where(
                WorkflowState.workflow_id.in_(workflow_ids)
            )
            result = await session.execute(workflow_stmt)
            
            await session.commit()
            
            count = result.rowcount
            logger.info(f"清理了 {count} 个旧工作流数据（{days}天前）")
            return count
            
        except Exception as e:
            await session.rollback()
            logger.error(f"删除旧知识库条目失败: {str(e)}")
            return 0
    
    @staticmethod
    async def generate_embedding(text: str) -> List[float]:
        """
        生成文本的嵌入向量
        
        Args:
            text: 要生成嵌入的文本
            
        Returns:
            嵌入向量（浮点数组）
        """
        try:
            # 获取嵌入模型
            from utils.embedding import get_embedding_model
            embedding_model = get_embedding_model()
            
            # 生成嵌入向量
            embedding = await embedding_model.embed_async(text)
            
            if not isinstance(embedding, list) or not embedding:
                raise ValueError("生成的嵌入向量无效")
            
            return embedding
            
        except Exception as e:
            logger.error(f"生成嵌入向量失败: {str(e)}")
            # 返回默认向量
            return [0.0] * 1536  # 假设使用1536维向量


# 全局实例
db_utils = DatabaseUtils()
